from typing import List, Optional, Set, Tuple
from tree_sitter import Node, Tree

from chernc.parser.rust_parser import parse_rust
from chernc.treediff.gumtree.matchers.mapping import Mapping
from chernc.treediff.gumtree.tree.tree import TreeLike
from chernc.treediff.gumtree.matchers.mapping_store import MappingStore
from chernc.treediff.gumtree.matchers.similarity_metrics import SimilarityMetrics
from chernc.treediff.gumtree.gen.generators.treesitter_tree import TreeSitterNode
from chernc.treediff.gumtree.tree.tree_utils import TreeUtils


def _gumtree_peek_max(lst: List[TreeLike]) -> int:
    max_heights = []
    for node in lst:
        max_heights.append(node.get_height())
    return max(max_heights)


def _gumtree_push_without_sort(t: TreeLike, lst: List[TreeLike]):
    # Push node t into list lst (ordered by decreasing height)
    lst.append(t)


def _gumtree_push(t: TreeLike, lst: List[TreeLike]):
    # Push node t into list lst (ordered by decreasing height)
    lst.append(t)
    lst.sort(key=lambda node: node.get_height(), reverse=True)


def _gumtree_pop(lst: List[TreeLike]) -> List[TreeLike]:
    # Pop all nodes with the greatest height from the list lst
    max_height = _gumtree_peek_max(lst)
    nodes_to_remove = [node for node in lst if node.get_height() == max_height]

    # Remove those nodes from lst
    for node in nodes_to_remove:
        lst.remove(node)

    return nodes_to_remove


def _gumtree_open(t: TreeLike, lst: List[TreeLike]):
    # Insert all children of node t into list lst
    for child in t.get_children():
        _gumtree_push_without_sort(child, lst)
    lst.sort(key=lambda node: node.get_height(), reverse=True)


def _gumtree_dice(t1: Optional[TreeLike], t2: Optional[TreeLike], mapping: MappingStore[TreeLike]) -> float:
    """
    计算两个树节点（t1 和 t2）之间的 Dice 相似度。
    :param t1: 源树节点
    :param t2: 目标树节点
    :param mapping: 存储映射关系的 MappingStore
    :return: 计算得到的 Dice 相似度
    """
    if t1 is None or t2 is None:
        return 0.0
    return SimilarityMetrics.dice_similarity(t1, t2, mapping)


def gumtree_topdown(source: TreeLike, destination: TreeLike, min_height: int = 2):
    src_height: List[TreeLike] = []
    dst_height: List[TreeLike] = []
    candidate_mapping: MappingStore[TreeLike] = MappingStore(source, destination) # A
    mapping: MappingStore[TreeLike] = MappingStore(source, destination) # M

    src_height.append(source)
    dst_height.append(destination)

    while min(_gumtree_peek_max(src_height), _gumtree_peek_max(dst_height)) > min_height:
        if _gumtree_peek_max(src_height) != _gumtree_peek_max(dst_height):
            if _gumtree_peek_max(src_height) > _gumtree_peek_max(dst_height):
                for t in _gumtree_pop(src_height):
                    _gumtree_open(t, src_height)
            else:
                for t in _gumtree_pop(dst_height):
                    _gumtree_open(t, dst_height)
        else:
            h1 = _gumtree_pop(src_height)
            h2 = _gumtree_pop(dst_height)
            for t1 in h1:
                for t2 in h2:
                    if t1.is_isomorphic_to(t2):
                        if any(t1.is_isomorphic_to(tx) and tx != t2 for tx in destination.get_descendants()) \
                            or any(t2.is_isomorphic_to(tx) and tx != t1 for tx in source.get_descendants()):
                            candidate_mapping.add_mapping(t1, t2)
                        else:
                            # Add pairs of isomorphic nodes to mapping M
                            candidate_mapping.add_mapping(t1, t2)
                            # Add all pairs of isomorphic nodes of s(t1) and s(t2) to M
                            mapping.add_mapping_recursively(t1, t2)
            # Open new children for nodes in H1 and H2 that aren't in the mapping or candidate list
            for t1 in h1:
                if not candidate_mapping.is_src_mapped(t1):
                    _gumtree_open(t1, src_height)
            for t2 in h2:
                if not candidate_mapping.is_dst_mapped(t2):
                    _gumtree_open(t2, dst_height)

    # Sort candidate mappings based on the Dice coefficient between parent nodes
    candidate_mapping_set: Set[Mapping[TreeLike]] = candidate_mapping.as_set()
    candidate_mapping_list: List[Mapping[TreeLike]] = list(candidate_mapping_set)
    candidate_mapping_list.sort(key=lambda pair: _gumtree_dice(pair.src.get_parent(), pair.dst.get_parent(), mapping))

    # 处理候选映射中的每一对
    while len(candidate_mapping_list) > 0:
        t1t2 = candidate_mapping_list.pop(0)
        t1, t2 = t1t2.src, t1t2.dst
        mapping.add_mapping_recursively(t1, t2)

        # 清除映射关系中的不必要项
        removed_indices = []
        for i in range(len(candidate_mapping_list)):
            if candidate_mapping_list[i].src == t1 or candidate_mapping_list[i].dst == t1:
                removed_indices.append(i)
        
        removed_indices = sorted(removed_indices, reverse=True)
        for i in removed_indices:
            candidate_mapping_list.pop(i)

    return mapping


def gumtree_bottomup(source: TreeLike, destination: TreeLike, mapping: MappingStore[TreeLike], min_dice: float=0.5, max_size: int=100):
    """
    Bottom-up algorithm for tree matching phase.
    :param source: Source tree
    :param destination: Destination tree
    :param mapping: Mapping store from the top-down phase
    :param min_dice: Minimum Dice similarity threshold for considering mappings
    :param max_size: Maximum size limit for nodes considered for refinement
    :return: Updated mapping store
    """
    
    # Perform post-order traversal of source tree
    post_order_source = TreeUtils.post_order(source)

    # Iterate through the nodes of the source tree in post-order
    for t1 in post_order_source:
        # Check if t1 is not matched and has matched children
        if not mapping.is_src_mapped(t1) and any(mapping.is_src_mapped(c) for c in t1.get_children()):
            # Get the candidate mapping for t1 in the destination tree
            t2 = mapping.get_dst_for_src(t1)

            # If a valid candidate t2 is found and the Dice similarity is greater than the threshold
            if t2 is not None and _gumtree_dice(t1, t2, mapping) > min_dice:
                # Add the mapping (t1, t2) to the result
                mapping.add_mapping(t1, t2)

                # If the maximum size constraint is not violated, perform refinement
                if max(len(t1.get_children()), len(t2.get_children())) < max_size:
                    # Call opt (optimization) function to refine the mapping
                    refined_mappings = opt(t1, t2)
                    
                    # Add refined mappings to the result
                    for ta, tb in refined_mappings:
                        if not mapping.is_src_mapped(ta) and not mapping.is_dst_mapped(tb) and ta.get_label() == tb.get_label():
                            mapping.add_mapping(ta, tb)
    
    return mapping

def treediff(source: Tree, destination: Tree):
    source_root = TreeSitterNode(source.root_node)
    destination_root = TreeSitterNode(destination.root_node)
    mapping = gumtree_topdown(source_root, destination_root)
    return gumtree_bottomup(source_root, destination_root, mapping)


def print_tree(code, node, indent=0):
    # Recursively print the tree
    if not isinstance(node, Node):
        return
    print(" " * indent + f"{node.type} ({node.start_byte}-{node.end_byte}): {code[node.start_byte:node.end_byte]}")
    for child in node.children:
        print_tree(code, child, indent + 2)


if __name__ == "__main__":
    lhs = """
// CHERNC_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:3
fn swap(left: &mut i32, right: &mut i32) {
    let temp = *left;
    *left = *right;
    *right = temp;
}
// CHERNC_END_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:3

// CHERNC_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:10
fn bubble_sort(arry: &mut [i32], length: usize) {
    for i in 0..length-1 {
        for j in 0..length-i-1 {
            if arry[j] > arry[j+1] { // changed to use '>' for bubble sort
                let (left, right) = arry.split_at_mut(j + 1); // split the array into two parts
                swap(&mut left[j], &mut right[0]); // swap the elements
            }
        }
    }
}
// CHERNC_END_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:10

// CHERNC_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:25
fn main() {
    let mut input_array = [12, 43, 9, 13, 67, 98, 101, 89, 3, 35];
    println!("bubble sort the array");

    bubble_sort(&mut input_array, 10);

    // print the array
    print_array(&input_array, 10);
}
// CHERNC_END_MARK: /data/wangjun/github/Chern-C/examples/bubblesort/bubblesort.c:25
"""
    rhs = """
fn main() {
    let mut input_array = [12, 43, 9, 13, 67, 98, 101, 89, 3, 35];
    println!("bubble sort the array");

    bubble_sort(&mut input_array);
    print_array(&input_array, 10);
}
"""
    lhs_bytes = lhs.encode('utf-8')
    rhs_bytes = rhs.encode('utf-8')
    lhs_tree = parse_rust(lhs)
    rhs_tree = parse_rust(rhs)

    diff_result = treediff(lhs_tree, rhs_tree)
    for mapping in diff_result:
        print(lhs_bytes[mapping.src.get_pos(): mapping.src.get_end_pos()])
        print(rhs_bytes[mapping.dst.get_pos(): mapping.dst.get_end_pos()])

    # code = "bubble_sort(&mut input_array);"
    # rhs_tree = parse_rust(code)
    # print_tree(code, rhs_tree.root_node)
